import pickle

class ThreeDData:
    """3D data class: x, y, z lists of coordinates, value list of corresponding values."""
    x = []
    y = []
    z = []
    value = []
    def size(self):
        return len(self.x)
    def check_size(self):
        nx = len(self.x)
        ny = len(self.y)
        nz = len(self.z)
        nvalue = len(self.value)
        if ny != nx:
            return False
        if nz != nx:
            return False
        if nvalue != nx:
            return False
        return True
    def min_x(self):
        return min(self.x)
    def max_x(self):
        return max(self.x)
    def min_y(self):
        return min(self.y)
    def max_y(self):
        return max(self.y)
    def min_z(self):
        return min(self.z)
    def max_z(self):
        return max(self.z)
    def min_value(self):
        return min(self.value)
    def max_value(self):
        return max(self.value)
    def to_string(self):
        s = ''
        s += 'i : x y z : expression value\n'
        for i in range(self.size()):
            s += str(i) + ' : ' + str(self.x[i]) + ' ' + str(self.y[i]) + ' ' + str(self.z[i]) + ' : ' + str(self.value[i]) + '\n' 
        return s

class TwoDData:
    """2D data class: x, y lists of coordinates, value list of corresponding values."""
    x = []
    y = []
    value = []
    def size(self):
        return len(self.x)
    def check_size(self):
        nx = len(self.x)
        ny = len(self.y)
        nvalue = len(self.value)
        if ny != nx:
            return False
        if nvalue != nx:
            return False
        return True
        
def load_array(filename):
    """Load array values from pickle file."""
    print 'Unpickling ' + filename
    pkl_file = open(filename, 'rb')
    a = pickle.load(pkl_file)
    pkl_file.close()
    return a
    
def slice(data, value, plane='z'):
    """Slice 3D data to 2D by fixing plane coordinate to value.
    
    The return value will be a 3D object with the slice plane coordinate restricted to value."""
    n = data.size()
    sliced_data = ThreeDData()
    sliced_data.x = []
    sliced_data.y = []
    sliced_data.z = []
    sliced_data.value = []
    if plane == 'z':
        for i in range(n): 
            if data.z[i] == value:
                sliced_data.x.append(data.x[i])
                sliced_data.y.append(data.y[i])
                sliced_data.z.append(data.z[i])
                sliced_data.value.append(data.value[i])
    return sliced_data
       
def trim(data, keep=0.001):
    """Trim 3"""
    n = data.size()
    keep_size = int(n * keep)
    values = sorted(data.value)
    keep_above = values[n - keep_size]
    print 'min:', min(data.value), ', max:', max(data.value), ', keep_size:', keep_size, ', keep_above:', keep_above
    keep_values = values[(n - keep_size):]
    print 'keep_values:', keep_values
    ''' '''
    trimmed_data = ThreeDData()
    trimmed_data.x = []
    trimmed_data.y = []
    trimmed_data.z = []
    trimmed_data.value = []
    for i in range(n): 
        if data.value[i] >= keep_above:
            trimmed_data.x.append(data.x[i])
            trimmed_data.y.append(data.y[i])
            trimmed_data.z.append(data.z[i])
            trimmed_data.value.append(data.value[i])
    return trimmed_data
        
    
data = ThreeDData()

# unpickle arrays
data.x = load_array('3dDataX.pkl')
data.y = load_array('3dDataY.pkl')
data.z = load_array('3dDataZ.pkl')
data.value = load_array('3dDataValue.pkl')

#print data.x
#print data.y
#print data.z
#print data.value
print data.size()
print data.check_size()
trimmed_data = trim(data, 0.001)
print 'trimmed_data:\n', trimmed_data.to_string()

sliced_data = slice(trimmed_data, 42) 
print 'sliced_data:\n', sliced_data.to_string()
   
